# database.py
import pymysql
from dotenv import load_dotenv
import os
from typing import List, Dict, Any, Optional, Union, Tuple

# 加载.env文件
load_dotenv()


class DBConnection:
    """数据库连接上下文管理类（确保事务使用同一连接）"""

    def __init__(self):
        self.helper = MySQLHelper(
            host=os.getenv("host"),
            user=os.getenv("user"),
            password=os.getenv("password"),
            db=os.getenv("db"),
            port=int(os.getenv("port", "3306")),
            charset=os.getenv("charset", "utf8")
        )

    def __enter__(self):
        self.helper.begin_transaction()
        return self.helper

    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            if exc_type is not None:
                self.helper.rollback()
            else:
                self.helper.commit()
        finally:
            self.helper.close()


class MySQLHelper:
    def __init__(self, host: str, user: str, password: str, db: str,
                 port: int = 3306, charset: str = 'utf8'):
        self.host = host
        self.user = user
        self.password = password
        self.db = db
        self.port = port
        self.charset = charset
        self.connection = None
        self.cursor = None
        self._transaction_active = False

    # 连接管理方法
    def connect(self):
        """建立数据库连接"""
        self.close()  # 先关闭旧连接
        self.connection = pymysql.connect(
            host=self.host,
            user=self.user,
            password=self.password,
            db=self.db,
            port=self.port,
            charset=self.charset,
            cursorclass=pymysql.cursors.DictCursor,
            autocommit=False
        )
        self.cursor = self.connection.cursor()
        return self.connection

    def close(self):
        """关闭连接"""
        if self._transaction_active:
            self.rollback()

        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()
        self.cursor = None
        self.connection = None

    # 事务管理方法
    def begin_transaction(self):
        """开始事务"""
        if not self._transaction_active:
            self.connect()
            self.connection.begin()
            self._transaction_active = True

    def commit(self):
        """提交事务"""
        if self._transaction_active and self.connection:
            self.connection.commit()
            self._transaction_active = False

    def rollback(self):
        """回滚事务"""
        if self._transaction_active and self.connection:
            try:
                self.connection.rollback()
            except:
                pass
            finally:
                self._transaction_active = False

    # 查询方法
    def fetch_one(self, query: str, params: Tuple = ()) -> Optional[Dict[str, Any]]:
        """查询单条记录"""
        self.cursor.execute(query, params)
        return self.cursor.fetchone()

    def fetch_all(self, query: str, params: Tuple = ()) -> List[Dict[str, Any]]:
        """查询多条记录"""
        self.cursor.execute(query, params)
        return self.cursor.fetchall()

    # 执行方法
    def execute(self, query: str, params: Union[Tuple, List[Tuple]] = ()) -> int:
        """执行SQL语句"""
        if isinstance(params, list):
            return self.cursor.executemany(query, params)
        return self.cursor.execute(query, params)

    # 快捷操作方法
    def insert(self, table: str, data: Dict[str, Any]) -> int:
        """插入单条数据"""
        columns = ', '.join(data.keys())
        placeholders = ', '.join(['%s'] * len(data))
        query = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"
        self.execute(query, tuple(data.values()))
        return self.cursor.lastrowid

    def insert_many(self, table: str, data_list: List[Dict[str, Any]]) -> int:
        """
        批量插入数据
        返回影响的行数
        """
        if not data_list:
            return 0

        columns = ', '.join(data_list[0].keys())
        placeholders = ', '.join(['%s'] * len(data_list[0]))
        query = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"

        params = [tuple(item.values()) for item in data_list]
        return self.execute(query, params)

    def update(self, table: str, data: Dict[str, Any], condition: str, condition_params: Tuple = ()) -> int:
        """更新数据"""
        set_clause = ', '.join([f"{key} = %s" for key in data.keys()])
        query = f"UPDATE {table} SET {set_clause} WHERE {condition}"
        params = tuple(data.values()) + condition_params
        return self.execute(query, params)

    def delete(self, table: str, condition: str, params: Tuple = ()) -> int:
        """删除数据"""
        query = f"DELETE FROM {table} WHERE {condition}"
        return self.execute(query, params)

    def execute_in_transaction(self, operations: List[Dict]):
        """在事务中执行多个操作"""
        self.begin_transaction()
        try:
            result = None
            for op in operations:
                try:
                    if op['type'] == 'insert':
                        result = self.insert(op['table'], op['data'])
                    elif op['type'] == 'update':
                        result = self.update(
                            op['table'],
                            op['data'],
                            op.get('condition', '1=1'),
                            op.get('condition_params', ())
                        )
                    elif op['type'] == 'delete':
                        result = self.delete(
                            op['table'],
                            op.get('condition', '1=1'),
                            op.get('condition_params', ())
                        )
                    elif op['type'] == 'execute':
                        result = self.execute(op['sql'], op.get('params', ()))
                except Exception as e:
                    print(f"操作失败: {op}，错误: {e}")
                    raise

            self.commit()
            return result
        except Exception as e:
            self.rollback()
            print(f"事务回滚！错误: {e}")
            raise
        finally:
            self.close()

# 快捷操作函数
def query_one(sql: str, params: Tuple = ()) -> Optional[Dict[str, Any]]:
    """快捷查询单条记录"""
    with DBConnection() as db:
        return db.fetch_one(sql, params)


def query_all(sql: str, params: Tuple = ()) -> List[Dict[str, Any]]:
    """快捷查询多条记录"""
    with DBConnection() as db:
        return db.fetch_all(sql, params)


def execute(sql: str, params: Union[Tuple, List[Tuple]] = ()) -> int:
    """快捷执行SQL语句"""
    with DBConnection() as db:
        return db.execute(sql, params)


def insert(table: str, data: Dict[str, Any]) -> int:
    """快捷插入单条数据"""
    with DBConnection() as db:
        return db.insert(table, data)



def insert_many(table: str, data_list: List[Dict[str, Any]]) -> int:
    """快捷批量插入数据"""
    with DBConnection() as db:
        return db.insert_many(table, data_list)


def update(table: str, data: Dict[str, Any], condition: str, condition_params: Tuple = ()) -> int:
    """快捷更新数据"""
    with DBConnection() as db:
        return db.update(table, data, condition, condition_params)


def delete(table: str, condition: str, params: Tuple = ()) -> int:
    """快捷删除数据"""
    with DBConnection() as db:
        return db.delete(table, condition, params)

def begin():
    """开始事务"""
    with DBConnection() as db:
        db.begin_transaction()

def commit():
    """提交事务"""
    with DBConnection() as db:
        db.commit()

def rollback():
    """回滚事务"""
    with DBConnection() as db:
        db.rollback()

def execute_transaction(operations: List[Dict]):
    """
    在事务中执行多个操作
    :param operations: 操作列表，格式同MySQLHelper.execute_in_transaction
    :return: 最后一个操作的结果
    """
    with DBConnection() as db:
        return db.execute_in_transaction(operations)